#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ADB辅助工具类
提供ADB命令的封装和错误处理
"""

import subprocess
import os
import time
import json
from typing import List, Dict, Optional, Tuple

class ADBNotFoundError(Exception):
    pass

class ADBHelper:
    def __init__(self, adb_path=None):
        # 优先使用传入的adb_path，否则使用环境变量
        if adb_path:
            self.adb_path = adb_path
        else:
            # 兼容旧用法
            import main
            self.adb_path = getattr(main, 'ADB_BIN', 'adb')
        self.devices = []
        self.current_device = None
    
    def _check_adb_exists(self):
        if not os.path.exists(self.adb_path):
            raise ADBNotFoundError(f"未找到ADB工具: {self.adb_path}")
    
    def check_adb_available(self) -> bool:
        """检查ADB是否可用"""
        try:
            self._check_adb_exists()
            result = subprocess.run([self.adb_path, 'version'], 
                                  capture_output=True, text=True, encoding='utf-8')
            return result.returncode == 0
        except (FileNotFoundError, ADBNotFoundError):
            return False
    
    def get_devices(self) -> List[str]:
        """获取连接的设备列表"""
        try:
            self._check_adb_exists()
            result = subprocess.run([self.adb_path, 'devices'], 
                                  capture_output=True, text=True, encoding='utf-8')
            lines = result.stdout.strip().split('\n')[1:]  # 跳过标题行
            
            devices = []
            for line in lines:
                if line.strip() and '\t' in line:
                    device_id, status = line.split('\t')
                    if status == 'device':
                        devices.append(device_id)
            
            self.devices = devices
            return devices
        except (FileNotFoundError, ADBNotFoundError):
            raise ADBNotFoundError(f"未找到ADB工具: {self.adb_path}")
        except Exception as e:
            print(f"获取设备列表失败: {e}")
            return []
    
    def set_device(self, device_id: str) -> bool:
        """设置当前设备"""
        if device_id in self.devices:
            self.current_device = device_id
            return True
        return False
    
    def execute_command(self, command: List[str], timeout: int = 30) -> Tuple[bool, str, str]:
        """执行ADB命令"""
        self._check_adb_exists()
        try:
            if self.current_device:
                full_command = [self.adb_path, '-s', self.current_device] + command
            else:
                full_command = [self.adb_path] + command
            
            result = subprocess.run(full_command, 
                                  capture_output=True, 
                                  text=True, 
                                  encoding='utf-8',
                                  timeout=timeout)
            
            return result.returncode == 0, result.stdout, result.stderr
        except subprocess.TimeoutExpired:
            return False, "", "命令执行超时"
        except Exception as e:
            return False, "", str(e)
    
    def capture_screenshot(self, output_path: str = "screenshot.png") -> bool:
        """获取设备截图"""
        # 获取截图
        success, stdout, stderr = self.execute_command(['shell', 'screencap', '-p', '/sdcard/screenshot.png'])
        if not success:
            print(f"截图失败: {stderr}")
            return False
        
        # 拉取截图到本地
        success, stdout, stderr = self.execute_command(['pull', '/sdcard/screenshot.png', output_path])
        if not success:
            print(f"拉取截图失败: {stderr}")
            return False
        
        return True
    
    def get_ui_xml(self, output_path: str = "ui.xml") -> bool:
        """获取当前界面的XML布局"""
        # 获取XML
        success, stdout, stderr = self.execute_command(['shell', 'uiautomator', 'dump', '/sdcard/ui.xml'])
        if not success:
            print(f"获取XML失败: {stderr}")
            return False
        
        # 拉取XML到本地
        success, stdout, stderr = self.execute_command(['pull', '/sdcard/ui.xml', output_path])
        if not success:
            print(f"拉取XML失败: {stderr}")
            return False
        
        return True
    
    def click_element(self, x: int, y: int) -> bool:
        """点击指定坐标"""
        success, stdout, stderr = self.execute_command(['shell', 'input', 'tap', str(x), str(y)])
        return success
    
    def input_text(self, text: str) -> bool:
        """输入文本"""
        success, stdout, stderr = self.execute_command(['shell', 'input', 'text', text])
        return success
    
    def press_key(self, keycode: int) -> bool:
        """按键"""
        success, stdout, stderr = self.execute_command(['shell', 'input', 'keyevent', str(keycode)])
        return success
    
    def swipe(self, x1: int, y1: int, x2: int, y2: int, duration: int = 500) -> bool:
        """滑动"""
        success, stdout, stderr = self.execute_command([
            'shell', 'input', 'swipe', str(x1), str(y1), str(x2), str(y2), str(duration)
        ])
        return success
    
    def get_device_info(self) -> Dict[str, str]:
        """获取设备信息"""
        info = {}
        
        # 获取设备型号
        success, stdout, stderr = self.execute_command(['shell', 'getprop', 'ro.product.model'])
        if success:
            info['model'] = stdout.strip()
        
        # 获取Android版本
        success, stdout, stderr = self.execute_command(['shell', 'getprop', 'ro.build.version.release'])
        if success:
            info['android_version'] = stdout.strip()
        
        # 获取屏幕分辨率
        success, stdout, stderr = self.execute_command(['shell', 'wm', 'size'])
        if success:
            info['screen_size'] = stdout.strip()
        
        # 获取屏幕密度
        success, stdout, stderr = self.execute_command(['shell', 'wm', 'density'])
        if success:
            info['screen_density'] = stdout.strip()
        
        return info
    
    def install_apk(self, apk_path: str) -> bool:
        """安装APK"""
        success, stdout, stderr = self.execute_command(['install', apk_path])
        return success
    
    def uninstall_app(self, package_name: str) -> bool:
        """卸载应用"""
        success, stdout, stderr = self.execute_command(['uninstall', package_name])
        return success
    
    def get_current_package(self) -> str:
        """获取当前前台应用包名"""
        success, stdout, stderr = self.execute_command([
            'shell', 'dumpsys', 'activity', 'activities', '|', 'grep', 'mResumedActivity'
        ])
        if success and stdout:
            # 解析输出获取包名
            lines = stdout.strip().split('\n')
            for line in lines:
                if 'mResumedActivity' in line:
                    # 提取包名
                    import re
                    match = re.search(r'([a-zA-Z0-9_.]+)/[a-zA-Z0-9_.]+', line)
                    if match:
                        return match.group(1)
        return ""
    
    def get_current_activity(self) -> str:
        """获取当前前台Activity"""
        success, stdout, stderr = self.execute_command([
            'shell', 'dumpsys', 'activity', 'activities', '|', 'grep', 'mResumedActivity'
        ])
        if success and stdout:
            lines = stdout.strip().split('\n')
            for line in lines:
                if 'mResumedActivity' in line:
                    # 提取Activity名
                    import re
                    match = re.search(r'([a-zA-Z0-9_.]+/[a-zA-Z0-9_.]+)', line)
                    if match:
                        return match.group(1)
        return ""
    
    def is_screen_on(self) -> bool:
        """检查屏幕是否亮起"""
        success, stdout, stderr = self.execute_command(['shell', 'dumpsys', 'power', '|', 'grep', 'Display Power'])
        if success and stdout:
            return 'ON' in stdout
        return False
    
    def wake_up_screen(self) -> bool:
        """唤醒屏幕"""
        success, stdout, stderr = self.execute_command(['shell', 'input', 'keyevent', '26'])  # POWER键
        time.sleep(0.5)
        success2, stdout2, stderr2 = self.execute_command(['shell', 'input', 'keyevent', '224'])  # WAKE_UP
        return success and success2
    
    def lock_screen(self) -> bool:
        """锁屏"""
        success, stdout, stderr = self.execute_command(['shell', 'input', 'keyevent', '26'])  # POWER键
        return success
    
    def take_screenshot_with_timestamp(self, output_dir: str = ".") -> str:
        """带时间戳的截图"""
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        filename = f"screenshot_{timestamp}.png"
        output_path = os.path.join(output_dir, filename)
        
        if self.capture_screenshot(output_path):
            return output_path
        return ""
    
    def get_element_by_coordinates(self, x: int, y: int, xml_path: str = "ui.xml") -> Optional[Dict]:
        """根据坐标获取元素信息"""
        if not os.path.exists(xml_path):
            return None
        
        try:
            import xml.etree.ElementTree as ET
            tree = ET.parse(xml_path)
            root = tree.getroot()
            
            def find_element_at_position(element, target_x, target_y):
                bounds = element.attrib.get('bounds', '')
                if bounds:
                    import re
                    match = re.match(r'\[(\d+),(\d+)\]\[(\d+),(\d+)\]', bounds)
                    if match:
                        x1, y1, x2, y2 = map(int, match.groups())
                        if x1 <= target_x <= x2 and y1 <= target_y <= y2:
                            return element
                
                for child in element:
                    result = find_element_at_position(child, target_x, target_y)
                    if result:
                        return result
                
                return None
            
            element = find_element_at_position(root, x, y)
            if element:
                return element.attrib
            return None
            
        except Exception as e:
            print(f"解析XML失败: {e}")
            return None 